# encoding:utf/8
from mmdet.apis import inference_detector, init_detector
import json
import os
from tqdm import tqdm


def result_from_dir():
    index = {1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6}
    # build the model from a config file and a checkpoint file
    model = init_detector(config_path, model_path, device='cuda:0')
    pics = os.listdir(pic_path)
    results = []
    num = 0
    for im in tqdm(pics):
        num += 1
        img = os.path.join(pic_path, im)
        result_ = inference_detector(model, img)

        # 从1开始
        for i, boxes in enumerate(result_, 1):
            if len(boxes):
                defect_label = index[i]
                for box in boxes:
                    d = {}
                    d["name"] = im
                    d['category'] = defect_label
                    d['bbox'] = [round(float(i), 2) for i in box[0:4]]
                    d['score'] = float(box[4])
                    results.append(d)

    with open(json_out_path, 'w') as fp:
        json.dump(results, fp)


if __name__ == "__main__":
    model_path = '/data/lzy/work_dir/c_rcnn_r50_fpn_2x_coco_RCCP_640_1280_add_ratio_add_dcn_add_size4_one2four_addmjx_no_dataaug/latest.pth'
    config_path = '/data/lzy/work_dir/c_rcnn_r50_fpn_2x_coco_RCCP_640_1280_add_ratio_add_dcn_add_size4_one2four_addmjx_no_dataaug/one2four_add_mjx.py'
    json_out_path = '/data/lzy/work_dir/c_rcnn_r50_fpn_2x_coco_RCCP_640_1280_add_ratio_add_dcn_add_size4_one2four_addmjx_no_dataaug/result_bak.json'
    pic_path = '/data/lzy/tile_round1_testA_20201231/testA_imgs/'
    result_from_dir()

    
